import numpy as np
import sympy as sp
import sys

sys.path.append('../pde')
sys.path.append('../fem')


from fealpy.mesh import MeshFactory
from fealpy.mesh import TetrahedronMesh
from linear_elasticity_model3D import GenLinearElasticitymodel3D
from fealpy.pde.linear_elasticity_model import QiModel3d, PolyModel3d
from LinearElasticityFEMModel import LinearElasticityFEMModel
from fealpy.tools.show import showmultirate


import matplotlib.pyplot as plt


m = int(sys.argv[1])
p = int(sys.argv[2])
n = int(sys.argv[3])
s = int(sys.argv[4])

E = 1e3
nu = 0.3
lam = E*nu/((1+nu)*(1-2*nu))
mu = E/(2*(1+nu))


pi = sp.pi
sin = sp.sin
cos = sp.cos
exp = sp.exp
ln = sp.ln
mf = MeshFactory()



#一般pde算例
x = sp.symbols('x0:3')
#u = [sin(2*pi*x[0]),sin(2*pi*x[1]),sin(2*pi*x[2])]
u = [exp(x[0]+x[1]),exp(x[0]-x[1]),x[1]+x[2]]
#u = [sin(2*pi*x[0])*sin(2*pi*x[1])*sin(2*pi*x[2]),x[0]*(1-x[0])*x[1]*(1-x[1])*x[2]*(1-x[2]),0]
#u = [exp(x[0]+x[1]+x[2]),x[1]*exp(x[0]+x[1]+x[2]),exp(x[0]+x[1]+x[2])]
#u = [exp(x[0]),x[1]*exp(x[2]),exp(x[1])]
#u = [x[0],x[1],x[2]]
lam = 1.0
mu = 0.5
if True:
    pde = GenLinearElasticitymodel3D(u,x,lam=lam,mu=mu,
                        Dirichletbd_n = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))', 
                        Dirichletbd_t0 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))', 
                        Dirichletbd_t1 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x1-1.0)<1e-12))|((abs(x2-0.0)<1e-12)|(abs(x2-1.0)<1e-12))')

elif False:
    pde = GenLinearElasticitymodel3D(u,x,lam=lam,mu=mu,
                        Dirichletbd_n = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)', 
                        Dirichletbd_t0 = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)', 
                        Dirichletbd_t1 = '((abs(x1-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-1.0)<1e-12))|(abs(x2-0.0)<1e-12)',
                        Neumannbd_nn='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)',
                        Neumannbd_nt0='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)',
                        Neumannbd_nt1='(abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12)')

elif False:
    pde = GenLinearElasticitymodel3D(u,x,lam=lam,mu=mu,
                        Dirichletbd_n  = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))', 
                        Dirichletbd_t0 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))', 
                        Dirichletbd_t1 = '((abs(x0-0.0)<1e-12)|(abs(x0-1.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))',
                        Neumannbd_nn='(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)',
                        Neumannbd_nt0='(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)',
                        Neumannbd_nt1='(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)')

else:
    pde = GenLinearElasticitymodel3D(u,x,lam=lam,mu=mu,
                        Dirichletbd_n  = '((abs(x0-0.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))', 
                        Dirichletbd_t0 = '((abs(x0-0.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))', 
                        Dirichletbd_t1 = '((abs(x0-0.0)<1e-12))|((abs(x1-0.0)<1e-12)|(abs(x2-0.0)<1e-12))',
                        Neumannbd_nn ='(abs(x0-1.0)<1e-12)|(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)',
                        Neumannbd_nt0='(abs(x0-1.0)<1e-12)|(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)',
                        Neumannbd_nt1='(abs(x0-1.0)<1e-12)|(abs(x1-1.0)<1e-12)|(abs(x2-1.0)<1e-12)')



#pde = QiModel3d(lam=1e8,mu=0.5)
#pde = PolyModel3d()

mesh = pde.init_mesh(n=n)


integrator = mesh.integrator(5)


showmesh = False
if showmesh:
    fig = plt.figure()
    axes = fig.gca()
    mesh.add_plot(axes)
    mesh.find_node(axes, showindex=True)
    mesh.find_cell(axes, showindex=True)
    plt.show()




if m == 1:
   
    
    maxit = 3
    errorType = ['$||\sigma - \sigma_h ||_{0}$',
                '$||div(\sigma - \sigma_h)||_{0}$',
                '$||u - u_h||_{0}$',
                '$||\sigma - \sigma_I ||_{0}$',
                '$||div(\sigma - \sigma_I)||_{0}$',
                '$||u - u_I||_{0}$',
                '$|| u_I - u_h||_{0}$'
                ]
    Ndof = np.zeros((maxit,))
    Internumber = np.zeros((maxit,))
    solve_time = np.zeros((maxit,))
    linear_operator_time = np.zeros((maxit,))
    errorMatrix = np.zeros((len(errorType), maxit), dtype=np.float)


    for i in range(maxit):
        fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
        fem.get_AA_and_bb()
        fem.solve()
        #fem.fast_solve()
        #fem.matlab_solve()
        #print('\n')
        fem.Finite_emelent_solution()
        tgdof = fem.tensorspace.number_of_global_dofs()
        vgdof = fem.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof
        Ndof[i] = gdof
        #Internumber[i] = fem.internumber
        #solve_time[i] = fem.solve_time
        #linear_operator_time[i] = fem.linear_operator_time
        errorMatrix[:, i] = fem.error()
        if i < maxit - 1:
            mesh.uniform_refine()
            
    print('Ndof:', Ndof)
    #print('Inter number:', Internumber)
    #print('solve time:',solve_time)
    #print('linear operator time:',linear_operator_time)
    print('error:', errorMatrix)
    showmultirate(plt, 0, Ndof, errorMatrix, errorType)
    plt.show()


elif m == 2:
    fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
    fem.Save_A_and_b()













